package cn.iocoder.yudao.module.engine.service.core.aichat.impl;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.date.DateUtil;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import cn.iocoder.yudao.module.engine.enums.AudioSendTaskStateEnum;
import cn.iocoder.yudao.module.engine.enums.EnumActionType;
import cn.iocoder.yudao.module.engine.enums.EnumEngineAipError;
import cn.iocoder.yudao.module.engine.exception.CancelLLMRequestException;
import cn.iocoder.yudao.module.engine.model.aichat.AIChatFuture;
import cn.iocoder.yudao.module.engine.model.aichat.AIChatReplayMessage;
import cn.iocoder.yudao.module.engine.model.aichat.AIChatReq;
import cn.iocoder.yudao.module.engine.model.config.EngineConfig;
import cn.iocoder.yudao.module.engine.model.context.AIMessageContext;
import cn.iocoder.yudao.module.engine.service.core.aichat.IAIChatService;
import cn.iocoder.yudao.module.infra.filestore.model.CreateFileRes;
import cn.iocoder.yudao.module.infra.filestore.FileStoreFactory;
import cn.iocoder.yudao.module.engine.service.core.aichat.splitter.TextSplitter;
import cn.iocoder.yudao.module.engine.service.core.aichat.splitter.TextSplitterFactory;
import cn.iocoder.yudao.module.infra.utils.EngineExecutorUtils;
import cn.iocoder.yudao.module.infra.llm.LLMServiceFactory;
import cn.iocoder.yudao.module.infra.llm.model.LLMRes;
import cn.iocoder.yudao.module.infra.tts.TTSServiceFactory;
import cn.iocoder.yudao.module.infra.util.GameMatchIdUtil;
import cn.iocoder.yudao.module.infra.util.TraceIdUtil;
import cn.iocoder.yudao.module.resource.api.IStoryApi;
import cn.iocoder.yudao.module.resource.api.model.CosplayDTO;
import cn.iocoder.yudao.module.resource.api.model.StoryDTO;
import cn.iocoder.yudao.module.resource.api.model.message.AudioDataDTO;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.*;

import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;

@Service
@Slf4j
public class AIChatServiceImpl implements IAIChatService {
    private ExecutorService llmExecutors;
    private ExecutorService ttsExecutors;

    @Resource
    private IStoryApi storyApi;
    @Resource
    private LLMServiceFactory llmServiceFactory;
    @Resource
    private TTSServiceFactory ttsService;
    @Resource
    private FileStoreFactory fileStoreFactory;
    @Resource
    private TextSplitterFactory textSplitterFactory;
    @Resource
    private EngineConfig engineConfig;

    @PostConstruct
    public void init() {
        final EngineConfig.ThreadPool threadPool = engineConfig.getThreadPool();
        llmExecutors = EngineExecutorUtils.newThreadPoolExecutor(threadPool.getLlmCorePoolSize(), "llm");
        ttsExecutors = EngineExecutorUtils.newThreadPoolExecutor(threadPool.getTtsCorePoolSize(), "tts");
    }

    @Override
    public AIChatFuture createMessage(AIChatReq req) {
        final AIChatFuture aiChatFuture = new AIChatFuture(req.getChatId(), req.getUserText());

        final String traceId = TraceIdUtil.getTraceId();
        final StoryDTO story = req.getStoryId() != null ? storyApi.getStory(req.getStoryId()) : null;

        final AIMessageContext messageCtx = new AIMessageContext(
                req.getGameMatchId(), traceId, req, story,
                aiChatFuture, textSplitterFactory.getService()
        );

        asyncStreamChat(messageCtx);
        return aiChatFuture;
    }

    @NotNull
    private Future<?> asyncStreamChat(AIMessageContext messageCtx) {
        final Semaphore userSemaphore = tryAcquireLLMByUserId(messageCtx);

        // 提交任务并转换为CompletableFuture
        final CompletableFuture<Void> llmFuture = CompletableFuture.supplyAsync(() -> {
            try {
                TraceIdUtil.setUserId(messageCtx.getReq().getUserId());
                TraceIdUtil.setTraceId(messageCtx.getTraceId());
                GameMatchIdUtil.setMatchId(messageCtx.getGameMatchId());
                streamChat(messageCtx);
            } catch (Exception e) {
                log.warn(StrUtil.format("AIChat异常, userText = {}",
                        messageCtx.getReq().getUserText()), e);
                messageCtx.completeWithError(messageCtx, e);
            } finally {
                userSemaphore.release();
                TraceIdUtil.removeUserId();
                TraceIdUtil.removeTraceId();
                GameMatchIdUtil.removeMatchId();
            }
            return null;
        }, llmExecutors);

        messageCtx.getAiChatFuture().setLlmFuture(llmFuture);
        llmFuture.whenComplete((result, throwable) -> {
            if (throwable != null) {
                if (throwable instanceof CancellationException) {
                    log.info("AIChat完成-LLM-取消, userText = {}",
                            messageCtx.getReq().getUserText());
                } else {
                    log.warn("AIChat完成-LLM-异常, userText = {}",
                            messageCtx.getReq().getUserText());
                }
            } else {
                log.info("AIChat完成-LLM-正常, userText = {}",
                        messageCtx.getReq().getUserText());
            }
        });

        return llmFuture;
    }

    private void streamChat(AIMessageContext messageCtx) {
        // 调用LLM服务，逐步生成文本
        messageCtx.setRequestLLMStartTime(System.currentTimeMillis());
        final String prompt = messageCtx.getStory() != null ? messageCtx.getStory().getPrompt() : messageCtx.getReq().getRole().getPrompt();
        llmServiceFactory.getService().streamChatResponse(
                messageCtx.getReq().getTone().getId(),
                prompt,
                messageCtx.getReq().getUserText(),
                messageCtx.getReq().getHistoryMessages(),
                llmRes -> {
                    if (messageCtx.getAiChatFuture().isCancelled()
                            || Thread.currentThread().isInterrupted()) {
                        log.info("取消LLM请求, userText = {}", messageCtx.getReq().getUserText());
                        throw new CancelLLMRequestException("取消LLM请求");
                    }

                    // setMessageId
                    if (StrUtil.isEmpty(messageCtx.getMessageId())) {
                        if (StrUtil.isNotEmpty(llmRes.getCacheMessageId())) {
                            messageCtx.setMessageId(llmRes.getCacheMessageId());
                        } else {
                            messageCtx.setMessageId(IdUtil.objectId());
                        }
                    }

                    messageCtx.setFromLLMCache(llmRes.isFromLLMCache());
                    messageCtx.setFromTTSCache(llmRes.isFromTTSCache());
                    final String replyText = llmRes.getReplyText();
                    final Boolean textDone = llmRes.isDone();

                    try {
                        // 推送文本消息
                        if (StrUtil.isNotEmpty(replyText)) {
                            sendReplyText(messageCtx, replyText);
                            messageCtx.getReplyTexts().add(replyText);
                            messageCtx.getSentence4AudioData().put(llmRes.getReplyText(), llmRes.getCacheAudios());
                        }

                        if (textDone) {
                            final String fullReplyText = String.join("", messageCtx.getReplyTexts());
                            messageCtx.setFullReplyText(fullReplyText);
                            messageCtx.setRequestLLMEndTime(System.currentTimeMillis());
                        }

                        if (messageCtx.getReq().getReturnAudio()) {
                            // 处理音频
                            handleReturnAudio(messageCtx, llmRes);
                        } else {
                            if (textDone) {
                                log.info("AIChat全部结束, userText = {}, fullReplyText = {}",
                                        messageCtx.getReq().getUserText(), messageCtx.getFullReplyText());
                                complete(messageCtx);
                            }
                        }
                    } catch (Exception e) {
                        log.warn(StrUtil.format("AIChat异常, userText = {}", messageCtx.getReq().getUserText()), e);
                        messageCtx.completeWithError(messageCtx, e);
                    }
                });
    }

    private void complete(AIMessageContext messageCtx) {
        messageCtx.complete(messageCtx);
    }

    // 处理音频
    private void handleReturnAudio(AIMessageContext messageCtx, LLMRes llmRes) {
        final Boolean textDone = llmRes.isDone();

        // 提取句子
        final List<String> sentenceList = getSentence(messageCtx.getTextSplitter(), llmRes);

        if (CollUtil.isNotEmpty(sentenceList)) {
            // 必须先增加Task，然后再设置isTextDone，否则会有并发问题
            final List<AIMessageContext.AudioTask> addedAudioTasks = new ArrayList<>();
            for (String sentence : sentenceList) {
                final AIMessageContext.AudioTask audioTask = newAudioTask(messageCtx, sentence);
                messageCtx.addAudioTaskQueue(audioTask);
                addedAudioTasks.add(audioTask);
            }

            if (textDone) {
                messageCtx.getTextDoneForAudio().set(true);
            }

            for (AIMessageContext.AudioTask audioTask : addedAudioTasks) {
                // 转换文本为音频并发送
                asyncConvertTextAndSendAudio(messageCtx, audioTask);
            }
        } else {
            if (textDone) {
                messageCtx.getTextDoneForAudio().set(true);
            }

            if (isAllAudioTaskDone(messageCtx)) {
                log.info("AIChat全部结束, sentenceList = {}", sentenceList);
                complete(messageCtx);
            }
        }
    }

    private void asyncConvertTextAndSendAudio(AIMessageContext messageCtx,
                                              AIMessageContext.AudioTask audioTask) {
        final String sentence = audioTask.getAudioData().getSentence();

        final CompletableFuture<String> ttsFuture = CompletableFuture.supplyAsync(() -> {
            Semaphore userSemaphore = null;
            try {
                userSemaphore = tryAcquireTtsByUserId(messageCtx);

                TraceIdUtil.setUserId(messageCtx.getReq().getUserId());
                TraceIdUtil.setTraceId(messageCtx.getTraceId());
                GameMatchIdUtil.setMatchId(messageCtx.getGameMatchId());

                convertTextToSpeech(messageCtx, audioTask);
                executeNextSendAudioTask(messageCtx);
            } catch (Exception e) {
                log.warn(StrUtil.format("TTS异常, taskId = {}, sentence = {}",
                        audioTask.getTaskId(), sentence), e);
                messageCtx.completeWithError(messageCtx, e);
            } finally {
                if (userSemaphore != null) {
                    userSemaphore.release();
                }
                log.info("TTS完成, taskId = {}, sentence = {}", audioTask.getTaskId(), sentence);
                TraceIdUtil.removeUserId();
                TraceIdUtil.removeTraceId();
                GameMatchIdUtil.removeMatchId();
            }
            return null;
        }, ttsExecutors);

        ttsFuture.whenComplete((result, throwable) -> {
            messageCtx.getAiChatFuture().removeTtsFuture(ttsFuture);
            if (throwable != null) {
                if (throwable instanceof CancellationException) {
                    log.info("AIChat完成-TTS-取消, tts.size = {}, sentence = {}, userText = {}",
                            messageCtx.getTtsFuturesSize(), sentence, messageCtx.getReq().getUserText());
                } else {
                    log.warn("AIChat完成-TTS-异常, tts.size = {}, sentence = {}, userText = {}",
                            messageCtx.getTtsFuturesSize(), sentence, messageCtx.getReq().getUserText());
                }
            } else {
                log.info("AIChat完成-TTS-正常, tts.size = {}, sentence = {}, userText = {}",
                        messageCtx.getTtsFuturesSize(), sentence, messageCtx.getReq().getUserText());
            }
        });
        messageCtx.getAiChatFuture().addTtsFuture(ttsFuture);
    }

    private void executeNextSendAudioTask(AIMessageContext messageCtx) {
        synchronized (messageCtx.getAudioTaskLock()) {
            final AIMessageContext.AudioTask waitingSendAudioTask = messageCtx.getFirstNeedSendTask();
            if (waitingSendAudioTask == null) {
                log.info("队首没有等待发送的音频任务");

                if (isAllAudioTaskDone(messageCtx)) {
                    log.info("TTS任务全部结束");
                    complete(messageCtx);
                }
                return;
            }

            // 发送音频
            try {
                waitingSendAudioTask.setSendState(AudioSendTaskStateEnum.SENDING);

                final String audioUrl = waitingSendAudioTask.getAudioData().getUrl();
                if (StrUtil.isEmpty(audioUrl)) {
                    log.warn("音频为空不发送, taskId = {}, sentence = {}",
                            waitingSendAudioTask.getTaskId(), waitingSendAudioTask.getAudioData().getSentence());
                } else {
                    saveAndSendAudioUrl(messageCtx, waitingSendAudioTask);
                }

                // 执行下一个任务
                executeNextSendAudioTask(messageCtx);
            } catch (Exception e) {
                log.warn(StrUtil.format("发送音频失败, taskId = {}, sentence = {}",
                        waitingSendAudioTask.getTaskId(), waitingSendAudioTask.getAudioData().getSentence()), e);
                messageCtx.completeWithError(messageCtx, e);
            }
        }
    }

    private boolean isAllAudioTaskDone(AIMessageContext messageCtx) {
        final boolean textDoneForAudio = messageCtx.getTextDoneForAudio().get();
        final int ttsNotDoneCount = messageCtx.getTtsNotDoneCount();
        final int remainingSendTaskCount = messageCtx.remainingSendTaskCount();
        log.info("判断TTS是否全部完成, textDoneForAudio = {}, ttsNotDoneCount = {}, remainingSendTaskCount = {}",
                textDoneForAudio, ttsNotDoneCount, remainingSendTaskCount);
        return textDoneForAudio && ttsNotDoneCount == 0 && remainingSendTaskCount == 0;
    }

    private void saveAndSendAudioUrl(AIMessageContext messageCtx, AIMessageContext.AudioTask task) throws IOException {
        // 发送音频
        try {
            final AIChatReplayMessage msg = new AIChatReplayMessage();
            msg.setAudioData(task.getAudioData());
            messageCtx.getAudioReplyCount().incrementAndGet();
            messageCtx.send(messageCtx, msg);
            log.info("发送音频成功, taskId = {}, sentence = {}, audioUrl = {}",
                    task.getTaskId(), task.getAudioData().getSentence(), task.getAudioData().getUrl());
        } finally {
            task.setSendState(AudioSendTaskStateEnum.DONE);
        }
    }

    @NotNull
    private static AIMessageContext.AudioTask newAudioTask(AIMessageContext messageCtx, String sentence) {
        final String taskId = IdUtil.objectId();
        final AudioDataDTO audioData = AudioDataDTO.builder()
                .toneId(messageCtx.getReq().getTone().getId())
                .id(null)
                .sentence(sentence)
                .url(null)
                .encryptKey(null)
                .audioBytes(null)
                .build();
        final AIMessageContext.AudioTask audioTask = new AIMessageContext.AudioTask(
                taskId, AudioSendTaskStateEnum.WAIT, false, audioData);
        return audioTask;
    }

    // 转换文本为语音
    private void convertTextToSpeech(AIMessageContext messageCtx,
                                     AIMessageContext.AudioTask audioTask) {
        try {
            final CosplayDTO role = messageCtx.getReq().getRole();
            final Integer toneId = messageCtx.getReq().getTone().getId();
            final String hash = messageCtx.getReq().getTone().getFileId().toString();
            final String sentence = audioTask.getAudioData().getSentence();

            // 读取缓存
            final List<AudioDataDTO> ttsAudios = messageCtx.getSentence4AudioData().get(sentence);
            if (CollUtil.isNotEmpty(ttsAudios)) {
                final Optional<AudioDataDTO> toneAudioOpt = ttsAudios.stream().filter(v -> v.getToneId().equals(toneId)).findFirst();
                if (toneAudioOpt.isPresent()) {
                    final AudioDataDTO cacheAudioData = toneAudioOpt.get();
                    audioTask.setAudioData(cacheAudioData);
                    log.info("TTS使用缓存, sentence = {}, url = {}", sentence, cacheAudioData.getUrl());
                    return;
                }
            }

            // TTS
            String ttsName = messageCtx.getReq().getTtsName();
            if (StrUtil.isEmpty(ttsName)) {
                ttsName = ttsService.getPrimaryName();
            }
            final String cloneVoiceId = messageCtx.getReq().getTone().getCloneVoiceIdByTssName(ttsName);
            final byte[] audioBytes = ttsService.getService(ttsName).convertTextToSpeech(
                    toneId, cloneVoiceId, role.getVoice(), role.getVoiceUrl(), hash, null,
                    messageCtx.getReq().getReturnAudioFormat(), sentence
            );
            final String audioId = IdUtil.objectId();
            audioTask.getAudioData().setId(audioId);
            audioTask.getAudioData().setAudioBytes(audioBytes);

            // 保存音频
            final String filePath = createFilePath(messageCtx, audioId);
            final CreateFileRes createFileRes = fileStoreFactory.getService().createFile(filePath, audioBytes);
            audioTask.getAudioData().setPath(filePath);
            audioTask.getAudioData().setUrl(createFileRes.getUrl());
            audioTask.getAudioData().setEncryptKey(createFileRes.getEncryptKey());
        } finally {
            audioTask.setTtsDone(true);
        }
    }

    private static String createFilePath(AIMessageContext messageCtx, String audioId) {
        final String returnAudioFormat = messageCtx.getReq().getReturnAudioFormat();
        final String date = DateUtil.format(new Date(), "yyyyMMdd");

        final String path;
        if (messageCtx.isFromLLMCache()) {
            path = StrUtil.format("tts/store/{}/{}.{}", date, audioId, returnAudioFormat);
        } else {
            path = StrUtil.format("tts/user/{}/{}.{}", date, audioId, returnAudioFormat);
        }
        return path;
    }

    private static List<String> getSentence(TextSplitter textSplitter, LLMRes llmRes) {
        if (llmRes.isCompleteSentence()) {
            // 完整句子
            return Collections.singletonList(llmRes.getReplyText());
        }

        final List<String> sentence = textSplitter.getSentence(llmRes.getReplyText());
        if (llmRes.isDone()) {
            final String remainingText = textSplitter.getRemainingText();
            if (StrUtil.isNotEmpty(remainingText)) {
                sentence.add(remainingText);
            }
        }
        log.info("提取句子, sentence = {}", sentence);
        return sentence;
    }

    private void sendReplyText(AIMessageContext messageCtx, String replyText) {
        final AIChatReplayMessage msg = new AIChatReplayMessage();
        msg.setText(replyText);
        try {
            messageCtx.send(messageCtx, msg);
            log.info("发送文本成功, replyText = {}", replyText);
        } catch (Exception e) {
            log.warn(StrUtil.format("发送文本失败, replyText = {}", replyText), e);
            messageCtx.completeWithError(messageCtx, e);
        }
    }

    private Semaphore tryAcquireLLMByUserId(AIMessageContext messageCtx) {
        // 获取或创建用户的信号量
        final Semaphore userSemaphore = messageCtx.getSemaphore(EnumActionType.LLM);

        // 尝试获取许可，如果没有可用许可会阻塞
        final boolean acquired;
        try {
            acquired = userSemaphore.tryAcquire(60, TimeUnit.SECONDS);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }

        if (!acquired) {
            log.error("LLM服务繁忙，请稍后再试, userId = {}, availablePermits = {}",
                    messageCtx.getReq().getUserId(), userSemaphore.availablePermits());
            throw exception(EnumEngineAipError.DEFAULT_ERROR, "LLM服务繁忙，请稍后再试");
        }
        return userSemaphore;
    }

    private Semaphore tryAcquireTtsByUserId(AIMessageContext messageCtx) throws InterruptedException {
        // 获取或创建用户的信号量
        final Semaphore userSemaphore = messageCtx.getSemaphore(EnumActionType.TTS);

        // 尝试获取许可，如果没有可用许可会阻塞
        final boolean acquired = userSemaphore.tryAcquire(60, TimeUnit.SECONDS);
        if (!acquired) {
            log.error("TTS服务繁忙，请稍后再试, userId = {}, availablePermits = {}",
                    messageCtx.getReq().getUserId(), userSemaphore.availablePermits());
            throw exception(EnumEngineAipError.DEFAULT_ERROR, "TTS服务繁忙，请稍后再试");
        }
        return userSemaphore;
    }

}
